//! # 限流存储实现
//!
//! 提供多种存储后端的限流数据存储实现，支持内存存储和Redis分布式存储。

use super::{RateLimitStorage, RateLimitResult, RateLimitError};
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, SystemTime};
use tokio::sync::RwLock;
use tracing::debug;

/// 内存存储项
#[derive(Debug, Clone)]
struct MemoryItem {
    value: u64,
    expires_at: SystemTime,
}

impl MemoryItem {
    fn new(value: u64, ttl: Duration) -> Self {
        Self {
            value,
            expires_at: SystemTime::now() + ttl,
        }
    }

    fn is_expired(&self) -> bool {
        SystemTime::now() > self.expires_at
    }

    fn update_value(&mut self, new_value: u64, ttl: Duration) {
        self.value = new_value;
        self.expires_at = SystemTime::now() + ttl;
    }
}

/// 内存限流存储
pub struct MemoryStorage {
    data: Arc<RwLock<HashMap<String, MemoryItem>>>,
}

impl MemoryStorage {
    pub fn new() -> Self {
        Self {
            data: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    pub async fn cleanup_expired(&self) -> usize {
        let mut data_guard = self.data.write().await;
        let initial_count = data_guard.len();
        data_guard.retain(|_, item| !item.is_expired());
        initial_count - data_guard.len()
    }
}

impl Default for MemoryStorage {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl RateLimitStorage for MemoryStorage {
    async fn get_counter(&self, key: &str) -> RateLimitResult<u64> {
        debug!("Getting counter for key: {}", key);
        let data_guard = self.data.read().await;
        
        if let Some(item) = data_guard.get(key) {
            if item.is_expired() {
                Ok(0)
            } else {
                Ok(item.value)
            }
        } else {
            Ok(0)
        }
    }

    async fn increment_counter(
        &self,
        key: &str,
        increment: u64,
        ttl: Duration,
    ) -> RateLimitResult<u64> {
        debug!("Incrementing counter for key: {} by {}", key, increment);
        let mut data_guard = self.data.write().await;
        
        let new_value = if let Some(item) = data_guard.get_mut(key) {
            if item.is_expired() {
                item.update_value(increment, ttl);
                increment
            } else {
                // 使用饱和加法防止溢出
                let new_val = item.value.saturating_add(increment);
                item.update_value(new_val, ttl);
                new_val
            }
        } else {
            data_guard.insert(key.to_string(), MemoryItem::new(increment, ttl));
            increment
        };
        
        Ok(new_value)
    }

    async fn set_counter(&self, key: &str, value: u64, ttl: Duration) -> RateLimitResult<()> {
        debug!("Setting counter for key: {} to {}", key, value);
        let mut data_guard = self.data.write().await;
        data_guard.insert(key.to_string(), MemoryItem::new(value, ttl));
        Ok(())
    }

    async fn delete_counter(&self, key: &str) -> RateLimitResult<()> {
        debug!("Deleting counter for key: {}", key);
        let mut data_guard = self.data.write().await;
        data_guard.remove(key);
        Ok(())
    }

    async fn get_ttl(&self, key: &str) -> RateLimitResult<Option<Duration>> {
        let data_guard = self.data.read().await;
        
        if let Some(item) = data_guard.get(key) {
            if item.is_expired() {
                Ok(None)
            } else {
                let now = SystemTime::now();
                if let Ok(remaining) = item.expires_at.duration_since(now) {
                    Ok(Some(remaining))
                } else {
                    Ok(None)
                }
            }
        } else {
            Ok(None)
        }
    }

    async fn batch_get(&self, keys: &[String]) -> RateLimitResult<Vec<u64>> {
        debug!("Batch getting {} keys", keys.len());
        let data_guard = self.data.read().await;
        let mut results = Vec::with_capacity(keys.len());
        
        for key in keys {
            if let Some(item) = data_guard.get(key) {
                if item.is_expired() {
                    results.push(0);
                } else {
                    results.push(item.value);
                }
            } else {
                results.push(0);
            }
        }
        
        Ok(results)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_memory_storage_basic_operations() {
        let storage = MemoryStorage::new();

        let value = storage.get_counter("test_key").await.unwrap();
        assert_eq!(value, 0);

        storage.set_counter("test_key", 10, Duration::from_secs(60)).await.unwrap();
        let value = storage.get_counter("test_key").await.unwrap();
        assert_eq!(value, 10);

        let new_value = storage.increment_counter("test_key", 5, Duration::from_secs(60)).await.unwrap();
        assert_eq!(new_value, 15);

        storage.delete_counter("test_key").await.unwrap();
        let value = storage.get_counter("test_key").await.unwrap();
        assert_eq!(value, 0);
    }

    #[tokio::test]
    async fn test_memory_storage_ttl() {
        let storage = MemoryStorage::new();

        storage.set_counter("ttl_key", 100, Duration::from_millis(100)).await.unwrap();
        
        let value = storage.get_counter("ttl_key").await.unwrap();
        assert_eq!(value, 100);

        tokio::time::sleep(Duration::from_millis(150)).await;

        let value = storage.get_counter("ttl_key").await.unwrap();
        assert_eq!(value, 0);
    }
}